import openai
import numpy as np
from src.utils.logger import get_logger

logger = get_logger(__name__)

class EmbeddingGenerator:
    def __init__(self, api_key=None):
        """Initialize with optional API key (can use env var instead)"""
        if api_key:
            openai.api_key = api_key
            
    def get_embedding(self, text, model="text-embedding-3-small"):
        """Convert text to vector embedding using OpenAI API"""
        try:
            response = openai.Embedding.create(
                input=text,
                model=model
            )
            embedding = response['data'][0]['embedding']
            return np.array(embedding, dtype=np.float32)
        except Exception as e:
            logger.error(f"Error generating embedding: {e}")
            return None
            
    def batch_get_embeddings(self, texts, model="text-embedding-3-small"):
        """Convert multiple texts to embeddings in a batch"""
        embeddings = []
        for text in texts:
            embedding = self.get_embedding(text, model)
            if embedding is not None:
                embeddings.append(embedding)
        return np.array(embeddings, dtype=np.float32)